/**
 * @file Yolo.h
 * @author RittoQin (1264145445@qq.com)
 * @brief MNN模型部署框架
 * @version 1.0
 * @date 2022/4/18
 *
 * @copyright Copyright SCUT RobotLab(c) 2022
 *
 */

#pragma once

#include <opencv2/core/mat.hpp>
#include <MNNDefine.h>
#include <MNNForwardType.h>
#include <Interpreter.hpp>
#include <pthread.h>
#include "Parameter.h"

/**
 * Yolo输出层属性
 */
typedef struct {
    std::string name;   // 层的名称
    int stride; // 步长
    std::vector<Anchor> anchors;    // 先验框尺寸
} YoloLayer;

/**
 * 边界框属性
 */
typedef struct BoxInfo {
    int x1; // xmin
    int y1; // ymin
    int x2; // xmax
    int y2; // ymax
    float score;    // 得分
    int label;  // 标签号
} BoxInfo;

/**
 * fast 自然指数函数
 * @param x 输入
 * @return 输出e^x
 */
inline float fast_exp(float x) {
    union {
        uint32_t i;
        float f;
    } v{};
    v.i = (1 << 23) * (1.4426950409 * x + 126.93490512f);   // 对uint32_t的操作 => flaot类型自然指数函数
    return v.f;
}

/**
 * sigmoid 函数
 * @param x 输入
 * @return 非线性概率
 */
inline float sigmoid(float x) {
    return 1.f / (1.f + fast_exp(-x));
}

class Yolo {
public:
    Yolo(); // 构造

    Yolo(const std::string &);  // 构造

//    inline void clear() {
//        if (__net != nullptr)
//        {
//            delete __net;
//            __net = nullptr;
//        }
//    }

    std::vector<BoxInfo> inference(cv::Mat &);  // 推理

    void drawBox(cv::Mat &, std::vector<BoxInfo> &);    // 绘制边界框

private:
    std::string __model_path;   // 模型路径
    std::vector<YoloLayer> __layers;    // Yolo层数据
    MNN::Interpreter* __net;    // 网络模型
    MNN::Session *__session;    // 输入输出节点内存分配符
    std::vector<cv::Scalar> __colors;
    std::mutex tensor_lock;

private:
    void setupEngine(); // 设置引擎

    void preProcess(cv::Mat &);   // 预处理

    std::vector<BoxInfo>
    postProcess(std::vector<YoloLayer> &, cv::Size, int, int, float, float, float);  // 后处理 -- 输出张量(边界框)提取

    void nms(std::vector<BoxInfo> &, float); // 极大值抑制

    void coordsTransform(std::vector<BoxInfo> &, int, cv::Size); // 坐标变换
};
